# app/rag/embedding/custom.py

from .base import EmbeddingProvider
import requests
import json
import numpy as np
import time
from typing import List
from ...core.logger import logger
import os
from ...core.config import settings


class CustomEmbedding(EmbeddingProvider):
    def __init__(
        self,
        model: str = "bge-m3",
        api_base: str = "http://direct.virtaicloud.com:45608/v1/embeddings",
        api_key: str = "",  # 有些自搭建服务可能不需要API密钥
    ):
        """
        初始化自定义 Embedding 服务
        :param model: 模型名，默认为 bge-m3
        :param api_base: 自定义 endpoint
        :param api_key: API密钥（可选）
        """
        self.model = model
        self.api_base = api_base
        self.api_key = api_key or os.getenv("CUSTOM_API_KEY", "")
        
        logger.info(f"🌐 Custom Embedding loaded: {model} at {api_base}")
    
    def _send_request(self, texts: List[str]) -> List[List[float]]:
        """发送HTTP请求获取嵌入向量"""
        # 过滤空文本
        filtered_texts = [text for text in texts if text and text.strip()]
        if not filtered_texts:
            return []
        
        # 准备请求头
        headers = {
            "Content-Type": "application/json"
        }
        
        if self.api_key and self.api_key != "EMPTY":
            headers["Authorization"] = f"Bearer {self.api_key}"
        
        # 准备请求体
        payload = {
            "model": self.model,
            "input": filtered_texts
        }
        
        try:
            # 发送请求，添加重试机制
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    response = requests.post(
                        self.api_base,
                        headers=headers,
                        json=payload,
                        timeout=120  # 使用增加后的超时时间
                    )
                    # 如果请求成功，跳出重试循环
                    if response.status_code == 200:
                        break
                    else:
                        logger.warning(f"嵌入请求尝试 {attempt + 1} 失败: HTTP {response.status_code}")
                        if attempt < max_retries - 1:  # 不是最后一次尝试，等待后重试
                            time.sleep(2 ** attempt)  # 指数退避
                        else:
                            # 所有重试都失败了
                            logger.error(f"❌ 嵌入请求最终失败: HTTP {response.status_code}, {response.text}")
                            # 返回默认零向量，确保Milvus插入不会失败
                            return [[0.0] * settings.EMBEDDING_DIM for _ in filtered_texts]
                except requests.exceptions.Timeout:
                    logger.warning(f"嵌入请求尝试 {attempt + 1} 超时")
                    if attempt < max_retries - 1:  # 不是最后一次尝试，等待后重试
                        time.sleep(2 ** attempt)  # 指数退避
                    if attempt == max_retries - 1:  # 最后一次尝试也失败了
                        logger.error("❌ 嵌入请求重试次数用尽，所有尝试都超时")
                        # 返回默认零向量，确保Milvus插入不会失败
                        return [[0.0] * settings.EMBEDDING_DIM for _ in filtered_texts]
                except Exception as e:
                    logger.warning(f"嵌入请求尝试 {attempt + 1} 异常: {str(e)}")
                    if attempt == max_retries - 1:  # 最后一次尝试也失败了
                        logger.error(f"❌ 嵌入请求最终异常: {str(e)}")
                        # 返回默认零向量，确保Milvus插入不会失败
                        return [[0.0] * settings.EMBEDDING_DIM for _ in filtered_texts]
            
            # 解析响应
            data = response.json()
            logger.debug(f"✅ 嵌入请求成功，返回 {len(data['data'])} 个向量")
            
            # 提取嵌入向量
            embeddings = [item["embedding"] for item in data["data"]]
            
            # 检查向量维度
            if embeddings:
                logger.info(f"✅ 嵌入向量维度: {len(embeddings[0])}维")
            
            return embeddings
        except Exception as e:
            logger.error(f"❌ 嵌入请求异常: {str(e)}")
            # 返回默认零向量，确保Milvus插入不会失败
            return [[0.0] * settings.EMBEDDING_DIM for _ in filtered_texts]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        # 调用自定义服务进行批量文档嵌入
        logger.debug(f"使用自定义服务嵌入 {len(texts)} 个文档")
        return self._send_request(texts)

    def embed_query(self, text: str) -> List[float]:
        # 调用自定义服务进行查询嵌入
        logger.debug(f"使用自定义服务嵌入查询: {text[:50]}...")
        embeddings = self._send_request([text])
        # 如果没有返回向量，返回默认零向量
        return embeddings[0] if embeddings else [0.0] * settings.EMBEDDING_DIM
        
    def get_embedding(self):
        """
        获取 embedding 实例
        （为了兼容原有代码）
        """
        return self